import numpy as np
import control.matlab

import pat.autopilot as ap


class Autopilot(ap.Autopilot):
    def __init__(self, _args=None, _dm=None, _time=None):
        ap.Autopilot.__init__(self)
        Xe, Ue = _dm.trim(); A, B = _dm.get_jacobian(Xe, Ue)
        for law in [DirectLaw, RateLaw, AttitudeLaw, VelocityLaw, PositionLaw]:
            self.add_mode(law(_dm, Xe, Ue, A, B))
        self.set_cur_mode(2)

class Step_z(ap.StepReference):
    def __init__(self, Xr, Ur): ap.StepReference.__init__(self, Xr, Ur, 2, "Step z", ampl=10.) # z


class DirectLaw(ap.LinearFeedback):
    def __init__(self, _dm, Xe, Ue, A, B):
      ap.LinearFeedback.__init__(self, _dm.sv_size, _dm.iv_size, Xe, Ue)
      self.set_gain(np.array([[0., 0., 0., 0., 0., 0.,  0., 0., 0. ,  0.,  0. , 0.],
                              [0., 0., 0., 0., 0., 0.,  0., 0., 0. ,  0.,  0. , 0.],
                              [0., 0., 0., 0., 0., 0.,  0., 0., 0. ,  0.,  0.,  0.],
                              [0., 0., 0., 0., 0., 0.,  0., 0., 0. ,  0.,  0. , 0.]]))
      self.name = "Direct"

class RateLaw(ap.LinearFeedback):
    def __init__(self, _dm, Xe, Ue, A, B):
      ap.LinearFeedback.__init__(self, _dm.sv_size, _dm.iv_size, Xe, Ue)
      self.set_gain(np.array([[0., 0., 0., 0., 0., 0.,  0., 0., 0. ,  0.,  0. , 0.],
                              [0., 0., 0., 0., 0., 0.,  0., 0., 0. ,  0.5, 0. , 0.],
                              [0., 0., 0., 0., 0., 0.,  0., 0., 0. ,  0.,  2.5, 0.],
                              [0., 0., 0., 0., 0., 0.,  0., 0., 0. ,  0.,  0. , 1.5]]))
      self.name = "Rate"

class AttitudeLaw(ap.LinearFeedback):
    def __init__(self, _dm, Xe, Ue, A, B):
      ap.LinearFeedback.__init__(self, _dm.sv_size, _dm.iv_size, Xe, Ue)
      self.set_gain(np.array([[0., 0., 0., 0., 0., 0.,  0., 0., 0. ,  0.,  0. , 0.],
                              [0., 0., 0., 0., 0., 0.,  2., 0., 0. ,  0.5, 0. , 0.],
                              [0., 0., 0., 0., 0., 0.,  0., 2., 0. ,  0.,  2.5, 0.],
                              [0., 0., 0., 0., 0.,-1.5, 0., 0., 0. ,  0.,  0. , 1.5]]))
      self.name = "Attitude"


class VelocityLaw(ap.LinearFeedback):
    def __init__(self, _dm, Xe, Ue, A, B):
        A1=A[_dm.sv_z:, _dm.sv_z:]
        B1=B[_dm.sv_z:,:]
        #                  z      v    alpha beta   phi  theta  psi     p       q       r
        self.Q = np.diag([0.01,  0.1,  0.5,   0.5,  0.5, 2.5,   2.,  0.0005, 0.0005, 0.0005])
        #                 eng  ail  ele  rud
        self.R = np.diag([0.3, 0.1, 0.7, 0.1])
        (K1, X1, E) = control.matlab.lqr(A1, B1, self.Q, self.R)
        K = np.zeros((_dm.p.input_nb, _dm.sv_size))
        K[:,_dm.sv_z:] = -K1
        refs = [Step_z(Xe, Ue)]
        ap.LinearFeedback.__init__(self, _dm.sv_size, _dm.iv_size, Xe, Ue, K, "Velocity", refs)

    def properties(self):
#        with np.printoptions(precision=3, suppress=True):
        np.set_printoptions(precision=3, suppress=True)
        props = "Q=diag({:s})\nR=diag({:s})\nK=\n{:s}".format(np.diag(self.Q), np.diag(self.R), self.K)
        return props
        



class PositionLaw(ap.ControlLaw):
    def __init__(self, _dm, Xe, Ue, A, B):
        ap.ControlLaw.__init__(self)
        self.name = "Position"



